import torch
import torch.nn as nn
import torchvision as tvs
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.transforms import functional as ttF
from torchvision.ops import roi_align
from PIL import Image
import skimage.io as io
import cv2
import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import spacy
import pickle
# print(np.__version__)
# exit(0)

coco = [{'supercategory': 'person', 'id': 1, 'name': 'person'},
 {'supercategory': 'vehicle', 'id': 2, 'name': 'bicycle'},
 {'supercategory': 'vehicle', 'id': 3, 'name': 'car'},
 {'supercategory': 'vehicle', 'id': 4, 'name': 'motorcycle'},
 {'supercategory': 'vehicle', 'id': 5, 'name': 'airplane'},
 {'supercategory': 'vehicle', 'id': 6, 'name': 'bus'},
 {'supercategory': 'vehicle', 'id': 7, 'name': 'train'},
 {'supercategory': 'vehicle', 'id': 8, 'name': 'truck'},
 {'supercategory': 'vehicle', 'id': 9, 'name': 'boat'},
 {'supercategory': 'outdoor', 'id': 10, 'name': 'traffic light'},
 {'supercategory': 'outdoor', 'id': 11, 'name': 'fire hydrant'},
 {'supercategory': 'outdoor', 'id': 13, 'name': 'stop sign'},
 {'supercategory': 'outdoor', 'id': 14, 'name': 'parking meter'},
 {'supercategory': 'outdoor', 'id': 15, 'name': 'bench'},
 {'supercategory': 'animal', 'id': 16, 'name': 'bird'},
 {'supercategory': 'animal', 'id': 17, 'name': 'cat'},
 {'supercategory': 'animal', 'id': 18, 'name': 'dog'},
 {'supercategory': 'animal', 'id': 19, 'name': 'horse'},
 {'supercategory': 'animal', 'id': 20, 'name': 'sheep'},
 {'supercategory': 'animal', 'id': 21, 'name': 'cow'},
 {'supercategory': 'animal', 'id': 22, 'name': 'elephant'},
 {'supercategory': 'animal', 'id': 23, 'name': 'bear'},
 {'supercategory': 'animal', 'id': 24, 'name': 'zebra'},
 {'supercategory': 'animal', 'id': 25, 'name': 'giraffe'},
 {'supercategory': 'accessory', 'id': 27, 'name': 'backpack'},
 {'supercategory': 'accessory', 'id': 28, 'name': 'umbrella'},
 {'supercategory': 'accessory', 'id': 31, 'name': 'handbag'},
 {'supercategory': 'accessory', 'id': 32, 'name': 'tie'},
 {'supercategory': 'accessory', 'id': 33, 'name': 'suitcase'},
 {'supercategory': 'sports', 'id': 34, 'name': 'frisbee'},
 {'supercategory': 'sports', 'id': 35, 'name': 'skis'},
 {'supercategory': 'sports', 'id': 36, 'name': 'snowboard'},
 {'supercategory': 'sports', 'id': 37, 'name': 'sports ball'},
 {'supercategory': 'sports', 'id': 38, 'name': 'kite'},
 {'supercategory': 'sports', 'id': 39, 'name': 'baseball bat'},
 {'supercategory': 'sports', 'id': 40, 'name': 'baseball glove'},
 {'supercategory': 'sports', 'id': 41, 'name': 'skateboard'},
 {'supercategory': 'sports', 'id': 42, 'name': 'surfboard'},
 {'supercategory': 'sports', 'id': 43, 'name': 'tennis racket'},
 {'supercategory': 'kitchen', 'id': 44, 'name': 'bottle'},
 {'supercategory': 'kitchen', 'id': 46, 'name': 'wine glass'},
 {'supercategory': 'kitchen', 'id': 47, 'name': 'cup'},
 {'supercategory': 'kitchen', 'id': 48, 'name': 'fork'},
 {'supercategory': 'kitchen', 'id': 49, 'name': 'knife'},
 {'supercategory': 'kitchen', 'id': 50, 'name': 'spoon'},
 {'supercategory': 'kitchen', 'id': 51, 'name': 'bowl'},
 {'supercategory': 'food', 'id': 52, 'name': 'banana'},
 {'supercategory': 'food', 'id': 53, 'name': 'apple'},
 {'supercategory': 'food', 'id': 54, 'name': 'sandwich'},
 {'supercategory': 'food', 'id': 55, 'name': 'orange'},
 {'supercategory': 'food', 'id': 56, 'name': 'broccoli'},
 {'supercategory': 'food', 'id': 57, 'name': 'carrot'},
 {'supercategory': 'food', 'id': 58, 'name': 'hot dog'},
 {'supercategory': 'food', 'id': 59, 'name': 'pizza'},
 {'supercategory': 'food', 'id': 60, 'name': 'donut'},
 {'supercategory': 'food', 'id': 61, 'name': 'cake'},
 {'supercategory': 'furniture', 'id': 62, 'name': 'chair'},
 {'supercategory': 'furniture', 'id': 63, 'name': 'couch'},
 {'supercategory': 'furniture', 'id': 64, 'name': 'potted plant'},
 {'supercategory': 'furniture', 'id': 65, 'name': 'bed'},
 {'supercategory': 'furniture', 'id': 67, 'name': 'dining table'},
 {'supercategory': 'furniture', 'id': 70, 'name': 'toilet'},
 {'supercategory': 'electronic', 'id': 72, 'name': 'tv'},
 {'supercategory': 'electronic', 'id': 73, 'name': 'laptop'},
 {'supercategory': 'electronic', 'id': 74, 'name': 'mouse'},
 {'supercategory': 'electronic', 'id': 75, 'name': 'remote'},
 {'supercategory': 'electronic', 'id': 76, 'name': 'keyboard'},
 {'supercategory': 'electronic', 'id': 77, 'name': 'cell phone'},
 {'supercategory': 'appliance', 'id': 78, 'name': 'microwave'},
 {'supercategory': 'appliance', 'id': 79, 'name': 'oven'},
 {'supercategory': 'appliance', 'id': 80, 'name': 'toaster'},
 {'supercategory': 'appliance', 'id': 81, 'name': 'sink'},
 {'supercategory': 'appliance', 'id': 82, 'name': 'refrigerator'},
 {'supercategory': 'indoor', 'id': 84, 'name': 'book'},
 {'supercategory': 'indoor', 'id': 85, 'name': 'clock'},
 {'supercategory': 'indoor', 'id': 86, 'name': 'vase'},
 {'supercategory': 'indoor', 'id': 87, 'name': 'scissors'},
 {'supercategory': 'indoor', 'id': 88, 'name': 'teddy bear'},
 {'supercategory': 'indoor', 'id': 89, 'name': 'hair drier'},
 {'supercategory': 'indoor', 'id': 90, 'name': 'toothbrush'}]

COCO_CLASSES = {'0':'_background'}
for item in coco:
    COCO_CLASSES[str(item['id'])] = item['name']
# print(len(COCO_CLASSES))
# exit(0)


images_dir = './data/coco/train2014'
device = 'cpu'

def preprocess(image_path):
    # img = io.imread(image_path)
    image = cv2.imread(image_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = ttF.to_tensor(image).to(device)
    # image = ttF.resize(image, [224, 224], antialias=True).to(device)
    # print(image.shape)
    # exit(0)
    return image


def visualize_detections(file_names, out_info):
    for i in range(len(file_names)):
        file_name = file_names[i]
        info = out_info[i]
        # print(info)
        image = cv2.imread(os.path.join(images_dir, file_name))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        fig, ax = plt.subplots(1, figsize=(5, 5))
        ax.imshow(image)
        print(info['labels'])
        if len(info['labels'].shape) == 0:
            x1, y1, x2, y2 = info['boxes']
            rect = patches.Rectangle((x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor='r', facecolor='none')
            ax.add_patch(rect)
            ax.text(x1, y1 - 10, f"{COCO_CLASSES[str(info['labels'])]}: {info['scores']:.2f}", color='white', fontsize=6, bbox=dict(facecolor='blue', alpha=0.5))
        else:
            for box, label, score in zip(info['boxes'], info['labels'], info['scores']):
                x1, y1, x2, y2 = box
                rect = patches.Rectangle(
                    (x1, y1), x2 - x1, y2 - y1,
                    linewidth=2, edgecolor='r', facecolor='none'
                )
                ax.add_patch(rect)
                ax.text(
                    x1, y1 - 10,
                    f"{COCO_CLASSES[str(label)]}: {score:.2f}",
                    color='white', fontsize=6, bbox=dict(facecolor='blue', alpha=0.5)
                )
        ax.axis('off')
        plt.show()




if __name__ == '__main__':
    model = tvs.models.detection.fasterrcnn_resnet50_fpn(weights=tvs.models.detection.FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
    model.eval()
    model.to(device)

    images = []
    file_names = []
    num = 0
    for file_name in os.listdir(images_dir):
        file_names.append(file_name)
        img = preprocess(os.path.join(images_dir, file_name))
        images.append(img)
        break
        num += 1
        if num == 10:
            break

    out = model(images)
    # print(out)
    # exit(0)
    choose_out = []
    threshold = 0.7
    for item in out:
        boxes = item['boxes']
        labels = item['labels']
        scores = item['scores']
        choose_id = torch.argwhere(scores > threshold).squeeze()
        # print(choose_id)
        item['boxes'] = boxes[choose_id].cpu().detach().numpy()
        item['labels'] = labels[choose_id].cpu().detach().numpy()
        item['scores'] = scores[choose_id].cpu().detach().numpy()
        choose_out.append(item)
    # # print(len(choose_out))
    # visualize_detections(file_names, choose_out)





    # 假设 model.backbone 是特征提取部分
    features = model.backbone(images[0].unsqueeze(0))['0']  # 获取特征图

    # 对每个目标框提取 ROI 特征
    roi_features = roi_align(
        features, 
        [torch.tensor(choose_out[0]['boxes'], dtype=torch.float32)], 
        output_size=(7, 7)  # 特征图大小
    )
    print("ROI 特征形状:", roi_features.shape)  # [N_roi, C, 7, 7]



  
    # # 加载预训练模型（英文）
    # """使用前要python -m spacy download zh_core_web_sm安装模型"""
    # """英文: en_core_web_sm(小模型)、en_core_web_trf(大模型,更准)   
    # 中文:zh_core_web_sm(需先安装:pip install zh_core_web_sm)"""
    # nlp = spacy.load("en_core_web_sm")  # 中文用 "zh_core_web_sm"

    # # 待分析的文本
    # with open('./data/dbpedia/all_data.pkl', 'rb') as f:
    #     text = pickle.load(f)
    
    # # 实体识别
    # doc = nlp(text)

    # # 打印识别结果
    # for ent in doc.ents:
    #     print(f"实体: {ent.text}, 类型: {ent.label_}")

    # # 输出示例：
    # # 实体: Apple, 类型: ORG
    # # 实体: U.K., 类型: GPE
    # # 实体: $1 billion, 类型: MONEY
    # # 实体: Steve Jobs, 类型: PERSON
    # # 实体: Apple, 类型: ORG
    # 实体: 1976, 类型: DATE